/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * @brief load data instruction ut for ascend910B1
 *
 */
#include <gtest/gtest.h>
#include "kernel_operator.h"
#include "lib/matmul/tiling.h"
#include "impl/matmul/modules/matmul_param.h"
#include "impl/matmul/modules/cmatrix_buffer/matmul_l0c_buffer.h"

using namespace std;
using namespace AscendC;
using namespace matmul;

#include "matmul_module_test_def.h"

namespace {

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB>
class MatmulImpl
: MATMUL_IMPORT_MODULE(MatmulL0CBuffer, EnUnitFlag(MM_CFG)) {
    using VAR_PARAMS =
        typename MatmulParams<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, GetMatmulVersion(MM_CFG)>::PARAMS;

public:
    MatmulImpl() {
        InitVar();
    }

    VAR_PARAMS& GetVar() {
        return var;
    }

    void SetParamKIter(int k) {
        var.kIter_ = k;
    }

    void InitVar() {
        var.tiling_ = &tiling;
        var.tpipe_ = &pipe;
        var.baseMN_ = 1024;
    }

private:
    TCubeTiling tiling;
    TPipe pipe;
    VAR_PARAMS var;
};
}

class test_matmul_l0c_buffer : public testing::Test {
protected:
    void SetUp() {}
    void TearDown() {}

private:
    using L0cT = float;

    using A_TYPE = matmul::MatmulType<AscendC::TPosition::TSCM, CubeFormat::ND, half>;
    using B_TYPE = matmul::MatmulType<AscendC::TPosition::TSCM, CubeFormat::ND, half>;
    using C_TYPE = matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, L0cT>;
    using BIAS_TYPE = matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float>;

    MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_NORM, void> enUnitFlagMM;
    MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_MDL, void> disUnitFlagMM;
};

TEST_F(test_matmul_l0c_buffer, enable_unit_flag) {

    enUnitFlagMM.InitL0CBuffer<L0cT>();
    auto co1Local = enUnitFlagMM.LoadL0CTensor<L0cT>();
    auto co1Local1 = enUnitFlagMM.AllocL0CLocalTensor<L0cT>();
    enUnitFlagMM.FreeL0CLocalTensor<L0cT>(co1Local1);
    enUnitFlagMM.ResetL0CEventStatus();

    enUnitFlagMM.SetParamKIter(1);
    ASSERT_TRUE(enUnitFlagMM.IsL0CLastIter(0));

    ASSERT_EQ(enUnitFlagMM.GetMmadUnitFlagCtrl(true), 3);
    ASSERT_EQ(enUnitFlagMM.GetMmadUnitFlagCtrl(false), 2);

    FixpipeParams<L0cT> fixpipeParams;
    enUnitFlagMM.SetFixpipeUnitFlag(fixpipeParams);
    ASSERT_EQ(fixpipeParams.unitFlag, 3);
}

TEST_F(test_matmul_l0c_buffer, disable_unit_flag) {

    disUnitFlagMM.InitL0CBuffer<L0cT>();
    disUnitFlagMM.GetVar().cMatrix_ = disUnitFlagMM.LoadL0CTensor<L0cT>();
    auto co1Local1 = disUnitFlagMM.AllocL0CLocalTensor<L0cT>();
    disUnitFlagMM.FreeL0CLocalTensor<L0cT>(co1Local1);
    disUnitFlagMM.ResetL0CEventStatus();

    disUnitFlagMM.SetParamKIter(1);
    ASSERT_FALSE(disUnitFlagMM.IsL0CLastIter(0));

    ASSERT_EQ(disUnitFlagMM.GetMmadUnitFlagCtrl(true), 0);
    ASSERT_EQ(disUnitFlagMM.GetMmadUnitFlagCtrl(false), 0);

    FixpipeParams<L0cT> fixpipeParams;
    disUnitFlagMM.SetFixpipeUnitFlag(fixpipeParams);
    ASSERT_EQ(fixpipeParams.unitFlag, 0);
}